import os
import shutil
import argparse

def classify_and_move_images(source_dir):
    """
    分类、重命名并移动指定文件夹中的图片。

    - 结尾是 '_hazy_restored.png' 的图片移动到 'hazy_restored' 文件夹。
    - 结尾是 '_restored.png' 的图片移动到 'restored' 文件夹。
    - 移动后删除对应的后缀。
    """
    # 1. 检查源文件夹是否存在
    if not os.path.isdir(source_dir):
        print(f"错误：文件夹 '{source_dir}' 不存在。")
        return

    # 2. 定义目标文件夹的路径
    hazy_restored_dir = os.path.join(source_dir, 'hazy_restored')
    restored_dir = os.path.join(source_dir, 'restored')

    # 3. 如果目标文件夹不存在，则创建它们
    os.makedirs(hazy_restored_dir, exist_ok=True)
    os.makedirs(restored_dir, exist_ok=True)
    print(f"目标文件夹已准备就绪:\n- {hazy_restored_dir}\n- {restored_dir}\n")

    # 定义后缀
    hazy_suffix = '_hazy_restored.png'
    restored_suffix = '_restored.png'

    file_count_hazy = 0
    file_count_restored = 0

    # 4. 遍历源文件夹中的所有文件
    for filename in os.listdir(source_dir):
        # 构造文件的完整路径
        src_path = os.path.join(source_dir, filename)

        # 确保我们只处理文件，跳过文件夹
        if not os.path.isfile(src_path):
            continue

        # 5. 检查文件后缀并处理
        # 必须先检查更具体的 '_hazy_restored.png' 后缀
        if filename.endswith(hazy_suffix):
            # 计算新文件名 (删除后缀)
            new_filename = filename.replace('_hazy_restored', '')
            dst_path = os.path.join(hazy_restored_dir, new_filename)
            
            # 移动并重命名文件
            shutil.move(src_path, dst_path)
            print(f"移动: {filename}  ->  hazy_restored/{new_filename}")
            file_count_hazy += 1

        elif filename.endswith(restored_suffix):
            # 计算新文件名 (删除后缀)
            new_filename = filename.replace('_restored', '')
            dst_path = os.path.join(restored_dir, new_filename)

            # 移动并重命名文件
            shutil.move(src_path, dst_path)
            print(f"移动: {filename}  ->  restored/{new_filename}")
            file_count_restored += 1
    
    print("\n--- 操作完成 ---")
    print(f"总共移动了 {file_count_hazy} 个文件到 'hazy_restored' 文件夹。")
    print(f"总共移动了 {file_count_restored} 个文件到 'restored' 文件夹。")


def main():
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description="根据文件名后缀对图片进行分类、重命名和移动。")
    parser.add_argument(
        '--target_folder', 
        default='/root/code/FeMasr/result/outdoor',
        help='包含要处理的图片的目标文件夹路径。'
    )
    
    args = parser.parse_args()
    
    # 调用主函数
    classify_and_move_images(args.target_folder)

if __name__ == '__main__':
    main()

